// EG4PGAction.cc -- implementation of the class
// Created: Jan. 26, 2019, transcribed by Asia SUN, asia.rabbit@163.com
// Last modified: Sep. 29, 2024, by Asia SUN

#include <iostream>
#include <map>
#include <fstream>
#include <unistd.h>
#include <TH1F.h>
#include <TRandom3.h>
#include <TGraph.h>
#include "G4GeneralParticleSource.hh"
#include "G4ParticleDefinition.hh"
#include "G4SystemOfUnits.hh"
#include "G4PhysicalConstants.hh"
#include "G4ParticleGun.hh"
#include "G4IonTable.hh"
#include "EG4PGAction.h"
#include "EG4ParaManager.h"

static const double PI = 3.141592653589;

EG4PGAction::EG4PGAction(bool usingGun) : G4VUserPrimaryGeneratorAction(),
fGPS(new G4GeneralParticleSource()), fGun(new G4ParticleGun(1)), fTargetZ(0.),
fTargetHDZ(), fT(0.), fUsingGun(usingGun){
  if(fUsingGun){
    const Node &c = EG4ParaManager::Instance()->CtrlPara();
    fA = c["az"][0].as<int>();
    fZ = c["az"][1].as<int>();
    fEkPerU = c["energy"].as<double>(); // in MeV/u
    fTargetZ = c["targetZ"].as<double>(); // in mm
    fTargetHDZ = 0.5*c["targetD"].as<double>(); // in mm
    SetThetaHist(c["thetaHist"].as<string>());
  } // end if

  fRdm = new TRandom3();
  fRdm->SetSeed(rand());
} // end ctor

EG4PGAction::~EG4PGAction(){
  delete fRdm;
  delete fGPS;
  delete fGun;
} // end dtor

void EG4PGAction::GeneratePrimaries(G4Event *event){
  if(fUsingGun) GunShoot(event);
  else GPSShoot(event);
} // end member function GeneratePrimaries

// shoot using fGPS
void EG4PGAction::GPSShoot(G4Event *event){
  fGPS->GeneratePrimaryVertex(event);
} // end member function GPSShoot

// shoot using fGun
void EG4PGAction::GunShoot(G4Event *event){
  // get the particle
  ConstructIon(fZ, fA);

  // position and time
  fPos[0] = 0.;
  fPos[1] = 0.;
  fPos[2] = fRdm->Uniform(-fTargetHDZ, fTargetHDZ) + fTargetZ;
  fGun->SetParticlePosition(fPos);
  fGun->SetParticleTime(fT);

  // direction
  fPhi = fRdm->Uniform(-45.*deg, 45.*deg);
  fTheta = hTheta->GetRandom(fRdm);
  // momentum and energy
  const double mass = fParticle->GetPDGMass();
  fEk = fEkPerU*mass/amu_c2;
  fPP = std::sqrt(fEk*(fEk + 2.*mass));
  fP.setRThetaPhi(fPP, fTheta, fPhi);
  fGun->SetParticleEnergy(fEk);
  fGun->SetParticleMomentumDirection(fP);

  fGun->GeneratePrimaryVertex(event);
} // end member function GunShoot

void EG4PGAction::ConstructIon(int Z, int A){
  fParticle = G4IonTable::GetIonTable()->GetIon(Z, A, 0.);
  fGun->SetParticleDefinition(fParticle);
  fGun->SetParticleCharge(Z*eplus);
} // end member funcion ConstructIon

// read from file to fill hTheta
void EG4PGAction::SetThetaHist(const string &file){
  if(access(file.data(), F_OK)){
    std::cerr << "EG4PGAction::SetThetaDistr: file " << file << " does not exist";
    exit(1);
  } // end if
  std::map<double, double> dm;
  TGraph g; g.SetNameTitle("gTheta", "theta distribution");
  std::ifstream fin(file.data());
  double x, y;
  while(fin >> x >> y){
    x *= 0.001;
    // std::cout << "x: " << x << ", y: " << y << std::endl; getchar(); // DEBUG
    g.SetPoint(g.GetN(), x, y);
    dm[x] = y;
  } // end while
  const double xmin = dm.begin()->first, xmax = dm.rbegin()->first;
  const int n = 2000;
  const double maxTheta = 0.5*PI;
  double dx = maxTheta / n;
  hTheta = new TH1F("hTheta", "hTheta", n, 0., maxTheta);
  for(int i = 0; i < n; i++){
    x = dx*i;
    if(x < xmin || x > xmax) y = 0.;
    else y = g.Eval(x);
    hTheta->SetBinContent(i, y);
  } // end for over i
  g.Write("", TObject::kOverwrite);
  hTheta->Write("", TObject::kOverwrite);
} // end member function SetThetaHist
